package MACArray

import MYMAC._
import MYAccumulators._
import MYFIFO._

import chisel3._
import chisel3.util._

class MACArray(matrixSize:Int=5, dataWidth:Int=8) extends Module {
    val io = IO(new Bundle{
        val dataIn = Input(Vec(matrixSize, SInt(dataWidth.W)))//the data input to the leftmost column
        val switchw = Input(Vec(matrixSize, Bool()))          //the switchw input to the leftmost  clolumn
        val weightIn = Input(Vec(matrixSize, UInt((dataWidth*matrixSize).W)))//the weight input for the entire matrix
        //val weightWe = Input(Vec(matrixSize, Bool()))
        val weightWe = Input(Bool())                          //when active, start to programming the weight for preparation

        val dataOut = Output(Vec(matrixSize, SInt(32.W)))         //the output data for every column
    })

    private val MACs = Array.fill(matrixSize*matrixSize) {//the MAC matrix 
        Module(new MYMAC(matrixSize, dataWidth))
    }

    val weightWes = Wire(Vec(matrixSize, Bool()))//the enable signal for programming weight for every column
    val programming = RegInit(0.U(1.W))          //the control register for programming weight for every column
    val progstep = RegInit(0.U(log2Ceil(matrixSize).W))//the counter for programming weight
    val weightTags = Wire(Vec(matrixSize, UInt(log2Ceil(matrixSize).W)))//the tags input for the topmost row
    

    when(io.weightWe && programming===0.U(1.W)){//when io.weightWe is assert, start programming
        programming := 1.U(1.W)
    } .elsewhen(programming===1.U(1.W) && progstep === (matrixSize-1).U){//programming will keep assert for matrixSize cycles
        programming := 0.U(1.W)
    }

    when(programming===1.U(1.W)){//counting
        progstep := progstep + 1.U
    } .otherwise {
        progstep := 0.U
    }

    //used for drive the weightWe and the weightTag of the top row
    for(i <- 0 until matrixSize){//connect the weigh`tWes and weightTags for first row
        weightWes(i) := programming
        weightTags(i) := progstep
    }
    /*
    the two signals are used to make sure that the weightIn is propagated
    from the top to the bottom. When weightIn propagated at the bottom, each
    row's weightTag is equal to matrixSize-1, and each MAC can buffer the 
    weightIn at that cycle
    */

    //connect/////////////////////////////////////////////////////////////////////////////////////////////////////////////
    for(i <- 0 until matrixSize-1){//eg: i=0~matrixSize-2
        for(j <- 0 until matrixSize){//eg: j=0~matrixSize-1;  
        //connect weight and accin from top to bottom
            MACs((i+1)*matrixSize+j).io.macIO.weightIn := MACs(i*matrixSize +j).io.macIO.weightDelay
            MACs((i+1)*matrixSize+j).io.macIO.weightWe := MACs(i*matrixSize +j).io.macIO.weightWeDelay
            MACs((i+1)*matrixSize+j).io.macIO.accIn := MACs(i*matrixSize +j).io.macIO.outDelay
            MACs((i+1)*matrixSize+j).io.macIO.weightTag := MACs(i*matrixSize +j).io.macIO.weightTagDelay
        }
    }

    for (i <- 0 until matrixSize){//eg: i=0~matrixSize-1
        for(k <- 0 until matrixSize-1){//eg: k=0~matrixSize-2
            //connect datain and switchw from left to right
            MACs(i*matrixSize+k+1).io.macIO.dataIn := MACs(i*matrixSize+k).io.macIO.dataDelay
            MACs(i*matrixSize+k+1).io.macIO.switchw := MACs(i*matrixSize+k).io.macIO.switchDelay
        }
    }

    //the MACArray output
    for(i <- 0 until matrixSize){//eg i=0~matrixSize-1
        io.dataOut(i) := MACs((matrixSize-1)*matrixSize+i).io.macIO.outDelay
    }

    //the leftmost MACArray input
    for(i <- 0 until matrixSize){//eg: i=0~matrixSize-1
        MACs(i*matrixSize).io.macIO.dataIn := io.dataIn(i)
        MACs(i*matrixSize).io.macIO.switchw := io.switchw(i)
    }
    /*
    the io.switchw signal maybe active for one cycle; but it will propagate for leftmost to rightmost column;
    after matrixSize cycles, the weigthIn propagate at the bottom row, and the switchw also propagate at the rightmost
    column; then the weightIn will be buffered in the same buffer(buf1 or buf2) for every MAC in MACs
    */

    //the topmost MACArray input
    //step1: select the correct rowWeight of matrix, the weight of last row need to be send first;
    val cases = io.weightIn.zipWithIndex.map{case (rowWeight, index) => 
        ((matrixSize-index-1).U===progstep) -> rowWeight
    }
    val currentWeight = MuxCase(0.U((dataWidth*matrixSize).W), cases)//the currentWeight bit-width is dataWidth*matrixSize
    //step2: divide the selcted rowWeight to pieces
    val rowWeightVec = Wire(Vec(matrixSize, UInt(dataWidth.W)))
    for(i <- 0 until matrixSize){
        rowWeightVec(i) := currentWeight((i+1)*dataWidth-1, i*dataWidth)
    }
    //step3: connect weight and others
    for (i <- 0 until matrixSize){
        MACs(i).io.macIO.accIn := 0.S(32.W)
        MACs(i).io.macIO.weightWe := weightWes(i)
        //MACs(i).io.macIO.weightIn := io.weightIn(i)
        MACs(i).io.macIO.weightIn := rowWeightVec(i).asSInt//UInt convert to SInt
        MACs(i).io.macIO.weightTag := weightTags(i)
    }



}

object MACMain extends App {
  // println("Generating the MACArray hardware")
  // emitVerilog(new MACArray(3, 8), Array("--target-dir", "generated"))
  // println("Generating the Accums hardware")
  // emitVerilog(new MYAccumulators(4, 3), Array("--target-dir", "generated"))

  println("Generating the FIFO hardware")
  emitVerilog(new MYFIFO(14), Array("--target-dir", "generated"))
}